{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Surface restructuring analysis of LAMMPS NEB trajectory:\n",
    "## Preparation of VASP optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print('Please cite DOI: 10.1021/acs.jpcc.9b04863.')\n",
    "print('This script analyzes LAMMPS NEB calculations of interlayer surface restructuring events (as prepared via MD2NEB.py) and generates VASP POSCAR files for DFT optimization of each event.')\n",
    "print('This script is interactive and requires user input. Please read carefully before proceeding:\\n')\n",
    "\n",
    "print('Note 1: Intended only for interlayer surface restructuring events in FCC metals.')\n",
    "print('Note 2: Intended only for periodic slab models with vacuum along the z-direction.')\n",
    "print('Note 3: Intended only for VASP DFT optimizations (ionic relaxation for IS/FS; dimer method for TS).')\n",
    "print('Note 4: Must have used DFT unit cell in previous LAMMPS analysis (MD2NEB.py).\\n')\n",
    "\n",
    "print('Note 5: Requires a LAMMPS DATA file containing the equilibrated box dimensions.')\n",
    "print('Note 6: Requires a VASP POSCAR file containing the DFT-optimized lattice vectors of the slab model.')\n",
    "print('Note 7: Requires a file listing the event numbers chosen to be optimized, line-by-line.')\n",
    "print('Note 8: Requires NEB images in the form of LAMMPS DUMP file, named \"coords.final.$i\", 0 <= $i <= Nimg-1.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract unit cell & system information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwd = os.getcwd()\n",
    "\n",
    "# Load LAMMPS data file after NPT equilibration\n",
    "name = input('Enter the name of the LAMMPS DATA file containing the equilibrated box dimensions : ')\n",
    "file_eq = pwd + '/' + name\n",
    "log = open(file_eq,'r')\n",
    "log_lines = log.readlines()\n",
    "log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "for line_index, line in enumerate(log_lines):\n",
    "    if line:\n",
    "        for l in line:\n",
    "            \n",
    "            if l == 'xlo':\n",
    "                xlo = float(line[0])\n",
    "                xhi = float(line[1])\n",
    "                \n",
    "            elif l == 'ylo':\n",
    "                ylo = float(line[0])\n",
    "                yhi = float(line[1])\n",
    "            \n",
    "            elif l == 'zlo':\n",
    "                zlo = float(line[0])\n",
    "                zhi = float(line[1])\n",
    "            \n",
    "            elif l == 'xy':\n",
    "                xy = float(line[0])\n",
    "                xz = float(line[1])\n",
    "                yz = float(line[2])\n",
    "\n",
    "            elif l == 'Atoms':\n",
    "                start = line_index + 2\n",
    "            \n",
    "            elif l == 'Velocities':\n",
    "                end = line_index - 1\n",
    "\n",
    "lat_lmp = np.array([[xhi-xlo, 0.0, 0.0], [xy, yhi-ylo, 0.0], [xz, yz, zhi-zlo]])    # LAMMPS NVT lattice matrix\n",
    "\n",
    "for line in log_lines[start:end]:\n",
    "    line[:2] = np.array(line[:2]).astype(int)    # Atom ID, atom type\n",
    "    line[2:5] = np.array(line[2:5]).astype(float)    # x, y, z\n",
    "\n",
    "log.close()\n",
    "\n",
    "sys = input('Enter the system name : ')\n",
    "\n",
    "ls_element = input('Enter the list of the elements, separated by space, in the same order as LAMMPS atom types : ')\n",
    "\n",
    "ls_type = []    # List of atom types\n",
    "for line in log_lines[start:end]:\n",
    "    Type = line[1]\n",
    "    if Type not in ls_type:\n",
    "        ls_type.append(Type)\n",
    "N_element = len(ls_type)    # Number of elements\n",
    "\n",
    "N_type = []    # List of number of atoms of each atom type\n",
    "for atom_type in ls_type:\n",
    "    N = 0\n",
    "    for line in log_lines[start:end]:\n",
    "        Type = line[1]\n",
    "        if Type == atom_type:\n",
    "            N += 1\n",
    "    N_type.append(N)\n",
    "N_type = ' '.join(str(n) for n in N_type)\n",
    "\n",
    "# Load reference VASP POSCAR file\n",
    "name = input('Enter the name of the VASP POSCAR file containing the DFT-optimized lattice vectors of the slab model : ')\n",
    "poscar = pwd + '/' + name\n",
    "log = open(poscar,'r')\n",
    "log_lines = log.readlines()\n",
    "\n",
    "scale = log_lines[1]    # Lattice scaling factor\n",
    "\n",
    "a = log_lines[2]    # DFT-optimized lattice\n",
    "b = log_lines[3]\n",
    "c = log_lines[4]\n",
    "\n",
    "log.close()\n",
    "\n",
    "F = 'F F F'    # Selective dynamics\n",
    "T = 'T T T'\n",
    "\n",
    "fix = input('Selective dynamics? Yes or No : ')\n",
    "if fix == 'Yes':\n",
    "    print('The fixed atoms are assumed to be numbered consecutively (e.g. height-by-height). Please quit now if not so.')\n",
    "    header = '%s\\n%s%s%s%s%s\\n%s\\n%s\\n%s\\n'%(sys, scale, a, b, c, ls_element, N_type, 'Selective dynamics', 'Direct')\n",
    "    fix_range = input('Enter the IDs of the first and last atom to be fixed, separated by space : ')\n",
    "    fix_range = fix_range.split()\n",
    "    fix_range = np.array(fix_range).astype(int)\n",
    "else:\n",
    "    header = '%s\\n%s%s%s%s%s\\n%s\\n%s\\n'%(sys, scale, a, b, c, ls_element, N_type, 'Direct')\n",
    "\n",
    "print('Unit cell & system information extracted.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## List of events to be optimized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of events chosen to be optimized after visual inspection of the NEB results\n",
    "name = input('Enter the name of the file listing the event numbers chosen to be optimized, line-by-line : \\n Warning: All specified events must visually contain a single, well-defined event. Please quit now and perform visual inspection if not so.\\n')\n",
    "file_ev = pwd + '/' name\n",
    "ev = open(file_ev,'r')\n",
    "ev_lines = ev.readlines()\n",
    "ev_lines = [line.split() for line in ev_lines]\n",
    "\n",
    "ls = []\n",
    "for line in ev_lines:\n",
    "    ls.append(line[0])\n",
    "    \n",
    "ev.close()\n",
    "\n",
    "print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract IS, TS, FS - VASP POSCAR preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IS = initial state\n",
    "# TS = transition state\n",
    "# FS = final state\n",
    "\n",
    "print('Preparing VASP POSCAR files of IS, TS, FS in respective folders under ./dft/ev_***/ for each event...')\n",
    "\n",
    "for l in ls:\n",
    "    \n",
    "    path = pwd + '/ev_' + str(l)\n",
    "    \n",
    "    file_log = path + '/log.lammps'    # List of NEB distance & energy of the images\n",
    "    log = open(file_log,'r')\n",
    "    log_lines = log.readlines()\n",
    "    log_lines = np.array(log_lines)\n",
    "    N_line = log_lines.shape[0]\n",
    "    log_lines = [line.split() for line in log_lines]\n",
    "    \n",
    "    file_mep = path + '/mep.txt'    # MEP = minimum energy pathway\n",
    "    mep = open(file_mep,'w')\n",
    "    \n",
    "    os.chdir(path1)\n",
    "    N_img = os.popen('ls -1 coords.final.* | wc -l').readlines()    # Extract number of images\n",
    "    N_img = [line.split() for line in N_img]\n",
    "    N_img = np.array(N_img).astype(int)\n",
    "    N_img = int(N_img)\n",
    "    \n",
    "    for i in range(N_img,0,-1):    # Extract the converged NEB energy landscape (last entry)\n",
    "        out = log_lines[N_line-1][-2*i] + '\\t' + log_lines[N_line-1][-2*i+1] + '\\n'\n",
    "        mep.write(out)\n",
    "\n",
    "    mep.close()\n",
    "    log.close()\n",
    "    \n",
    "    mep = open(file_mep,'r')    # Scan MEP for IS, TS, FS\n",
    "    mep_lines = mep.readlines()\n",
    "    mep_lines = [line.split() for line in mep_lines]\n",
    "    mep_lines = np.array(mep_lines).astype(float)\n",
    "    \n",
    "    f = plt.figure()    # Save MEP plot as a PDF file\n",
    "    plt.plot(mep_lines, 'k.', markersize=10)\n",
    "    plt.plot(mep_lines, 'k--', markersize=1)\n",
    "    plt.ylabel('Energy (eV)')\n",
    "    plt.xlabel('Image number')\n",
    "    f.savefig('mep.pdf')\n",
    "    \n",
    "    for line_index, line in enumerate(mep_lines):        # TS = highest-energy image\n",
    "        if line[1] == np.nanmax(mep_lines[:,1]):\n",
    "            TS = line_index\n",
    "        \n",
    "    for i in range(TS,0,-1):    # Start at TS and climb down to the left to find IS\n",
    "        if mep_lines[i,1] - mep_lines[i-1,1] > 0:\n",
    "            if i > 1:\n",
    "                continue\n",
    "            else:\n",
    "                IS = i-1\n",
    "        else:\n",
    "            IS = i\n",
    "            break\n",
    "    \n",
    "    for i in range(TS,N_img-1):    # Start at TS and climb down to the right to find FS\n",
    "        if mep_lines[i,1] - mep_lines[i+1,1] > 0:\n",
    "            if i < N_img-2:\n",
    "                continue\n",
    "            else:\n",
    "                FS = i+1\n",
    "        else:\n",
    "            FS = i\n",
    "            break\n",
    "    \n",
    "    file_Nimg = path + '/Nimg.txt'    # List of image numbers corresponding to IS, TS, FS\n",
    "    Nimg = open(file_Nimg,'w')\n",
    "        \n",
    "    for i_index, i in enumerate([IS, TS, FS]):\n",
    "        for j_index, j in enumerate(['IS','TS','FS']):\n",
    "            if i_index == j_index:\n",
    "                os.mkdir(path + '/' + j)    # Create separate directories for IS, TS, FS\n",
    "                \n",
    "                out = j + '\\t' + str(i) + '\\n'\n",
    "                Nimg.write(out)\n",
    "                \n",
    "                file_poscar = path + '/' + j + '/POSCAR'    # VASP POSCAR files\n",
    "                poscar = open(file_poscar,'w')\n",
    "                poscar.write(header)\n",
    "                \n",
    "        file_img = path + '/coords.final.' + str(i)    # LAMMPS DUMP files of converged NEB images\n",
    "        img = open(file_img,'r')\n",
    "        img_lines = img.readlines()\n",
    "        img_lines = [line.split() for line in img_lines]\n",
    "        \n",
    "        for k in range(1,N_element+1):    # List atoms element-by-element\n",
    "            \n",
    "            for line_index, line in enumerate(img_lines):\n",
    "                if line_index > 8:\n",
    "                    \n",
    "                    line[0:2] = np.array(line[0:2]).astype(int)\n",
    "                    line[2:] = np.array(line[2:]).astype(float)\n",
    "                    \n",
    "                    ID = line[0]\n",
    "                    Type = line[1]\n",
    "\n",
    "                    if Type == k:\n",
    "                        \n",
    "                        x = line[2]\n",
    "                        y = line[3]\n",
    "                        z = line[4]\n",
    "                        r = np.array([x, y, z])    # Cartesian coordinates\n",
    "                        Rep = np.dot(r, np.linalg.inv(lat_lmp))    # Direct coordinates\n",
    "\n",
    "                        if fix == 'Yes':    # Selective dynamics\n",
    "                            if fix_range[0] <= ID <= fix_range[1]:\n",
    "                                coord = '%f %f %f %s\\n'%(Rep[0], Rep[1], Rep[2], F)\n",
    "                                poscar.write(coord)\n",
    "                            else:\n",
    "                                coord = '%f %f %f %s\\n'%(Rep[0], Rep[1], Rep[2], T)\n",
    "                                poscar.write(coord)\n",
    "                            \n",
    "                        else:\n",
    "                            coord = '%f %f %f\\n'%(Rep[0], Rep[1], Rep[2])\n",
    "                            poscar.write(coord)\n",
    "        \n",
    "        img.close()\n",
    "        poscar.close()\n",
    "    \n",
    "    Nimg.close()\n",
    "    mep.close()\n",
    "\n",
    "print('POSCAR generated > ./ev_***/{IS/,TS/,FS/}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
